#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IR 赋值语句节点定义
"""

import ast
from dataclasses import dataclass, field
from typing import Optional, Iterator, List, Union, TYPE_CHECKING
from ...symbol_table import Symbol, SymbolKind
from .base import IRStmt
from ..base import IRNode, get_default_ir_node
from ..expressions import (
    build_expr_from_ast,
    IRExpr,
    IRName,
    IRAttribute,
)
from ..expressions.call import IRCall
from ...symbol_table import TypeInfo


@dataclass
class IRAssign(IRStmt):
    """赋值语句"""

    # IR 节点表示 - 统一使用 target_exprs 存储所有目标
    target_exprs: List[IRExpr] = field(default_factory=list)  # 所有目标表达式

    # 赋值类型标记
    is_tuple_unpacking: bool = field(default=False)  # 是否是元组解包赋值

    value: IRNode = field(default_factory=get_default_ir_node)
    annotation: Optional[str] = field(default=None)
    node_type: str = "assign"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        # 所有目标表达式
        for target in self.target_exprs:
            yield target

        # 值节点
        yield self.value
        yield from super().get_child_nodes()

    def get_target_expr(self) -> Optional[IRExpr]:
        """获取第一个目标表达式（向后兼容接口）"""
        return self.target_exprs[0] if self.target_exprs else None

    def get_target_exprs(self) -> List[IRExpr]:
        """获取所有目标表达式"""
        return self.target_exprs

    def is_local_variable_assignment(self) -> bool:
        """判断是否是局部变量赋值

        Returns:
            True 如果是局部变量赋值，False 如果是属性赋值或其他类型
        """
        # 检查是否有任何局部变量符号
        return len(list(self.get_local_symbols())) > 0

    def get_local_symbols(self) -> Iterator["Symbol"]:
        """获取赋值语句中定义的所有局部变量符号

        处理所有类型的赋值：
        - 单目标赋值：x = 10
        - 多目标赋值：a = b = c = 10
        - 元组解包：x, y = 1, 2

        Returns:
            局部变量符号的迭代器
        """
        # 统一处理：从所有目标表达式中提取局部变量符号
        for target_expr in self.target_exprs:
            yield from self._extract_symbols_from_target(target_expr)

    def _extract_symbols_from_target(self, target_expr: IRExpr) -> Iterator["Symbol"]:
        """从目标表达式中递归提取局部变量符号

        Args:
            target_expr: 目标表达式

        Returns:
            局部变量符号的迭代器
        """
        if isinstance(target_expr, IRName):
            # 简单变量名：x = 10
            yield target_expr.get_symbol()
        elif hasattr(target_expr, "elements"):
            # 检查是否是元组或列表类型（有 elements 属性）
            # 元组或列表解包：(x, y) = ... 或 [x, y] = ...
            for element in getattr(target_expr, "elements"):
                yield from self._extract_symbols_from_target(element)
        # 注意：属性赋值（如 self.x = 10）和索引赋值（如 arr[0] = 10）
        # 不是局部变量定义，所以不在这里处理

    def build_from_ast(self, node) -> None:
        """从 AST 赋值节点构建 IR 赋值节点内容

        Args:
            node: AST 赋值节点 (ast.Assign 或 ast.AnnAssign)
        """
        super().build_from_ast(node)

        if isinstance(node, ast.AnnAssign):
            # 处理注解赋值
            self.build_from_ann_assign_ast(node)
        elif isinstance(node, ast.Assign):
            # 处理赋值语句
            if len(node.targets) == 1:
                # 单目标赋值或解包赋值
                target_ast = node.targets[0]
                if isinstance(target_ast, (ast.Tuple, ast.List)):
                    # 解包赋值：a, b = value
                    target_exprs = []
                    for elt in target_ast.elts:
                        target_node = build_expr_from_ast(elt, self)
                        if target_node and isinstance(target_node, IRExpr):
                            target_exprs.append(target_node)

                    if target_exprs:
                        self.target_exprs = target_exprs
                        self.is_tuple_unpacking = True  # 标记为元组解包
                else:
                    # 单目标赋值
                    # 普通单目标赋值（包括下标赋值）
                    target_node = build_expr_from_ast(target_ast, self)
                    if target_node and isinstance(target_node, IRExpr):
                        self.target_exprs = [target_node]  # 统一使用 target_exprs

            elif len(node.targets) > 1:
                # 多目标赋值：a = b = value
                target_exprs = []
                for target_ast in node.targets:
                    target_node = build_expr_from_ast(target_ast, self)
                    if target_node and isinstance(target_node, IRExpr):
                        target_exprs.append(target_node)

                if target_exprs:
                    self.target_exprs = target_exprs

            # 构建值表达式
            self.value = build_expr_from_ast(node.value, self)

    def build_from_ann_assign_ast(self, node: ast.AnnAssign) -> None:
        """从 AST 注解赋值节点构建 IR 赋值节点内容

        Args:
            node: AST 注解赋值节点
        """
        self.source_line = getattr(node, "lineno", 0)
        self.line = self.source_line

        # 处理目标变量
        target_node = build_expr_from_ast(node.target, self)
        if target_node and isinstance(target_node, IRExpr):
            self.target_exprs = [target_node]  # 统一使用 target_exprs
        else:
            # 无法处理的目标类型
            return

        # 处理类型注解
        if node.annotation:
            self.annotation = ast.unparse(node.annotation)

        # 处理初始值（如果有）
        if node.value:
            self.value = build_expr_from_ast(node.value, self)
        # 如果没有初始值，保持 value 为默认值（get_default_ir_node()）
        # 在代码生成阶段会根据类型生成合适的默认值

    def _handle_special_assignment(self, target: ast.AST, value: ast.AST) -> None:
        """处理特殊赋值情况（如索引赋值）

        Args:
            target: AST 目标节点
            value: AST 值节点
        """
        if isinstance(target, ast.Subscript):
            # 检查是否是切片赋值
            if isinstance(target.slice, ast.Slice):
                # 切片赋值：保留为 IRSubscript，在代码生成阶段处理
                # 构建目标和值的 IR 节点
                target_node = build_expr_from_ast(target, self)
                if target_node and isinstance(target_node, IRExpr):
                    self.target_exprs = [target_node]

                value_node = build_expr_from_ast(value, self)
                if value_node and isinstance(value_node, IRExpr):
                    self.value = value_node
                return

            # 处理索引赋值，如 dict[key] = value
            # 将其转换为 __setitem__ 方法调用
            # 构建 container.__setitem__(key, value) 调用
            container_expr = build_expr_from_ast(target.value, self)
            key_expr = build_expr_from_ast(target.slice, self)
            value_expr = build_expr_from_ast(value, self)

            # 创建 __setitem__ 方法调用
            setitem_attr = IRAttribute(parent=self)
            setitem_attr.value = container_expr
            setitem_attr.attr = "__setitem__"

            setitem_call = IRCall(parent=self)
            setitem_call.func = setitem_attr
            setitem_call.args = [key_expr, value_expr]

            # 将整个赋值转换为方法调用表达式
            self.value = setitem_call
            # 索引赋值没有左值目标，target_expr 保持为 None


@dataclass
class IRAugAssign(IRStmt):
    """增强赋值语句"""

    target: str = field(default="")
    op: str = field(default="")
    value: IRNode = field(default_factory=get_default_ir_node)
    node_type: str = "augassign"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.value
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.AugAssign) -> None:
        """从 AST 增强赋值节点构建 IR 增强赋值节点内容

        Args:
            node: AST 增强赋值节点
        """
        super().build_from_ast(node)

        if isinstance(node.target, ast.Name):
            self.target = node.target.id

            # 操作符映射
            op_map = {
                ast.Add: "+",
                ast.Sub: "-",
                ast.Mult: "*",
                ast.Div: "/",
                ast.Mod: "%",
                ast.Pow: "**",
                ast.LShift: "<<",
                ast.RShift: ">>",
                ast.BitOr: "|",
                ast.BitXor: "^",
                ast.BitAnd: "&",
                ast.FloorDiv: "//",
            }
            self.op = op_map.get(type(node.op), "unknown")

            self.value = build_expr_from_ast(node.value, self)
